package org.visionlibrary.image.filters.clustering;

import java.util.ArrayList;
import java.util.List;

import javax.media.jai.TiledImage;

import org.visionlibrary.image.acquisition.ImageFactory;
import org.visionlibrary.image.colormodel.HSVConverter;
import org.visionlibrary.image.colormodel.RGBConverter;
import org.visionlibrary.image.filters.clustering.proto.model.Cluster;
import org.visionlibrary.image.filters.clustering.proto.model.ColorSample;
import org.visionlibrary.image.filters.clustering.proto.model.CounterSample;
import org.visionlibrary.image.filters.clustering.proto.model.GraySample;
import org.visionlibrary.image.filters.clustering.proto.model.HSVCluster;
import org.visionlibrary.image.filters.clustering.proto.model.Sample;
import org.visionlibrary.image.geomtric.util.Utils;
import org.visionlibrary.image.model.AbstractFilter;


public class KMeansHSVAndGrayClusterization extends AbstractFilter {
	protected double hueQunatization = 12;
	protected double saturationQuantization = 0.125;
	protected double valueQuantization = 0.125;

	protected double kMax = 5;
	protected double tHMax = 25;
	protected double diffDistThres = 120;
	protected int iterations = 10;
	
	protected int hueQntLvl = 30;
	protected int satQntLvl = 8;
	protected int valQntLvl = 8;
	
	protected List<Cluster<Sample>> clusters = null;
	protected List<Sample> samples = null;
	
	public KMeansHSVAndGrayClusterization(double hueQunatization,
			double saturationQuantization, double valueQuantization,
			double kMax, double tHMax, int iterations) {
		super();
		this.hueQunatization = hueQunatization;
		this.saturationQuantization = saturationQuantization;
		this.valueQuantization = valueQuantization;
		this.kMax = kMax;
		this.tHMax = tHMax;
		this.iterations = iterations;
		
		Sample.setHueQuantization(hueQunatization);
		Sample.setSatQunatization(saturationQuantization);
		Sample.setValQuantization(valueQuantization);
		
		hueQntLvl = (int) Math.floor(360 / hueQunatization);
		satQntLvl = (int) Math.floor(1.0d / saturationQuantization);
		valQntLvl = (int) Math.floor(1.0d / valueQuantization);
	}

	@Override
	public TiledImage applyFilter(TiledImage src, TiledImage dest) {
		if (null == src)
			throw new NullPointerException("Source image is null.");

		if (null == dest)
			dest = (new ImageFactory()).createCompatibleImage(src);

		samples = new ArrayList<Sample>();
		
		int maxX = src.getWidth();
		int maxY = src.getHeight();

		for (int x = 0; x < maxX; x++)
			for (int y = 0; y < maxY; y++) {
				float[] values = new float[src.getNumBands()];
				for (int ch = 0; ch < src.getNumBands(); ch++) {
					int pix = src.getSample(x, y, ch);
					values[ch] = (float) Utils.norm(0, 255, 0, 1, pix);
				}
				float[] hsv = (new HSVConverter()).convertFromTYPE_RGB(values);
				int qHi = quantize(hsv[0], hueQunatization, hueQntLvl);
				int qSi = quantize(hsv[1], saturationQuantization, satQntLvl);
				int qVi = quantize(hsv[2], valueQuantization, valQntLvl);
				
				Sample sample = null;
				if(0 == qVi || 0 == qSi)
					sample = new GraySample(qVi);
				else
					sample = new ColorSample(qHi, qSi, qVi);
				
				if(samples.contains(sample))
					((CounterSample)samples.get(samples.indexOf(sample))).increment();
				else
					samples.add(sample);
			}
		
		double maxBinValue = Double.MIN_VALUE;
		Sample fstColCentroid = null;
		
		for(Sample sample : samples) {
			int value = ((CounterSample)sample).getCount();
			if (value > maxBinValue) {
				maxBinValue = value;
				fstColCentroid = sample;
			}
		}
		
		int centroidCount = 1;
		clusters = new ArrayList<Cluster<Sample>>();
		clusters.add(new HSVCluster(fstColCentroid));
		
		while (centroidCount < kMax) {
			double maxGlobalDist = Double.MIN_VALUE;
			Sample candidateCentroid = null; 

			for(Sample sample : samples) {
				double minLocalDist = Double.MAX_VALUE;
				for (Cluster<Sample> cluster : clusters) {
					Sample centroid = cluster.getCentroid();
					double dist = sample.distanceFrom(centroid);
					if (minLocalDist > dist)
						minLocalDist = dist;
				}
				if (maxGlobalDist < minLocalDist) {
					maxGlobalDist = minLocalDist;
					candidateCentroid = sample;
				}
			}

			if (!(maxGlobalDist >= tHMax))
				break;

			clusters.add(new HSVCluster(candidateCentroid));
			centroidCount++;
		}
		
		System.err.println("Clusters found: " + clusters.size());
		
		int i = 1;
		while (i <= iterations) {
			Cluster<Sample> bestMatch = null;
			for(Sample sample : samples) {
				double minDist = Double.MAX_VALUE;
				for (Cluster<Sample> cluster : clusters) {
					Sample centroid = cluster.getCentroid();
					double distance = sample.distanceFrom(centroid);
					if (minDist > distance) {
						bestMatch = cluster;
						minDist = distance;
					}
				}
				if (null != bestMatch)
					bestMatch.add(sample);
			}

			for (Cluster<Sample> cluster : clusters) 
				cluster.recalculate();

			if (++i <= iterations)
				for (Cluster<Sample> cluster : clusters)
					cluster.clear();
		}
		
		for (int x = 0; x < maxX; x++)
			for (int y = 0; y < maxY; y++) {
				float[] values = new float[src.getNumBands()];
				for (int ch = 0; ch < src.getNumBands(); ch++) {
					int pix = src.getSample(x, y, ch);
					values[ch] = (float) Utils.norm(0, 255, 0, 1, pix);
				}

				float[] hsv = (new HSVConverter()).convertFromTYPE_RGB(values);
				hsv = change(hsv, hueQntLvl, satQntLvl, valQntLvl);
				values = (new RGBConverter()).convertFromTYPE_HSV(hsv);

				for (int ch = 0; ch < src.getNumBands(); ch++) {
					int pix = (int) Utils.norm(0, 1, 0, 255, values[ch]);
					dest.setSample(x, y, ch, pix);
				}
			}

		return dest;
	}
	
	private float[] change(float[] hsv, int hueQntLvl, int satQntLvl, int valQntLvl) {
		float[] result = new float[hsv.length];
		int qHi = quantize(hsv[0], hueQunatization, hueQntLvl);
		int qSi = quantize(hsv[1], saturationQuantization, satQntLvl);
		int qVi = quantize(hsv[2], valueQuantization, valQntLvl);

		Sample smpl = null;
		if(0 == qVi || 0 == qSi)
			smpl = new GraySample(qVi);
		else
			smpl = new ColorSample(qHi, qSi, qVi);
		
		for (Cluster<Sample> cluster : clusters) {
			List<Sample> clusterSamples = cluster.getSamples();
			for (Sample sample : clusterSamples) {
				if (sample.equals(smpl)) {
					Sample centroid = cluster.getCentroid();
					result[0] = (float) (centroid.getHue() * hueQunatization);
					result[1] = (float) (centroid.getSat() * saturationQuantization);
					result[2] = (float) (centroid.getVal() * valueQuantization);

					return result;
				}
			}
		}
		throw new IllegalArgumentException(smpl.toString());
	}
	
	private int quantize(double what, double level, int max) {
		int value = (int) Math.floor((what / level));
		if(max <= value)
			value = max-1;
		
		return value;
	}
}	
